// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <list>
#include <memory>
#include <mutex>

namespace kumo::pollux::process {
    /// A registry for keeping static thread local objects of type T.  Similar to
    /// melon::ThreadLocal but a little bit more efficient in terms of performance
    /// and memory usage, because we do not support thread local with lexical scope.
    ///
    /// NOTE: only one instance of ThreadLocalRegistry<T> can be created with each
    /// T.
    template<typename T>
    class ThreadLocalRegistry {
    public:
        class Reference;

        /// Access values from all threads.  Takes a global lock and should be used
        /// with caution.
        template<typename F>
        void forAllValues(F f) {
            std::lock_guard<std::mutex> entriesLock(entriesMutex_);
            for (auto &entry: entries_) {
                std::lock_guard<std::mutex> lk(entry->mutex);
                f(entry->value);
            }
        }

    private:
        struct Entry {
            std::mutex mutex;
            T value;
        };

        std::list<std::unique_ptr<Entry> > entries_;
        std::mutex entriesMutex_;
    };

    /// Reference to one thread local value.  Should be stored in thread local
    /// memory.
    template<typename T>
    class ThreadLocalRegistry<T>::Reference {
    public:
        explicit Reference(const std::shared_ptr<ThreadLocalRegistry> &registry)
            : registry_(registry) {
            auto entry = std::make_unique<Entry>();
            std::lock_guard<std::mutex> lk(registry_->entriesMutex_);
            iterator_ =
                    registry_->entries_.insert(registry_->entries_.end(), std::move(entry));
        }

        ~Reference() {
            std::lock_guard<std::mutex> lk(registry_->entriesMutex_);
            registry_->entries_.erase(iterator_);
        }

        /// Obtain the thread local value and process it with the functor `f'.
        template<typename F>
        auto withValue(F f) {
            auto *entry = iterator_->get();
            std::lock_guard<std::mutex> lk(entry->mutex);
            return f(entry->value);
        }

    private:
        std::shared_ptr<ThreadLocalRegistry> const registry_;
        typename std::list<std::unique_ptr<Entry> >::iterator iterator_;
    };
} // namespace kumo::pollux::process
